#!/usr/bin/env python3
"""
Import NetworkX graph to Neo4j database
"""

import pickle
import sys
from neo4j import GraphDatabase
import networkx as nx
from tqdm import tqdm

# Neo4j connection settings
URI = "bolt://localhost:7687"
USERNAME = "neo4j"
PASSWORD = "neo4j123"  # We'll set this password

class Neo4jImporter:
    def __init__(self, uri, username, password):
        self.driver = GraphDatabase.driver(uri, auth=(username, password))

    def close(self):
        self.driver.close()

    def clear_database(self):
        """Clear existing data in the database"""
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            print("✅ Cleared existing database")

    def create_indexes(self):
        """Create indexes for better performance"""
        with self.driver.session() as session:
            # Create index on node id
            session.run("CREATE INDEX node_id IF NOT EXISTS FOR (n:Node) ON (n.id)")
            # Create index on node type
            session.run("CREATE INDEX node_type IF NOT EXISTS FOR (n:Node) ON (n.type)")
            print("✅ Created indexes")

    def import_nodes(self, G):
        """Import all nodes from NetworkX graph"""
        with self.driver.session() as session:
            nodes = list(G.nodes(data=True))

            # Batch import nodes
            batch_size = 1000
            for i in tqdm(range(0, len(nodes), batch_size), desc="Importing nodes"):
                batch = nodes[i:i+batch_size]

                node_data = []
                for node_id, attrs in batch:
                    node_type = attrs.get('type', 'unknown')
                    # Clean the node ID for display
                    label = node_id.split('/')[-1] if '/' in node_id else node_id

                    node_data.append({
                        'id': node_id,
                        'label': label,
                        'type': node_type,
                        'in_degree': G.in_degree(node_id),
                        'out_degree': G.out_degree(node_id),
                        'degree': G.degree(node_id)
                    })

                # Cypher query to create nodes
                query = """
                UNWIND $nodes AS node
                CREATE (n:Node {
                    id: node.id,
                    label: node.label,
                    type: node.type,
                    in_degree: node.in_degree,
                    out_degree: node.out_degree,
                    degree: node.degree
                })
                """

                session.run(query, nodes=node_data)

            print(f"✅ Imported {len(nodes)} nodes")

    def import_edges(self, G):
        """Import all edges from NetworkX graph"""
        with self.driver.session() as session:
            edges = list(G.edges(data=True))

            # Batch import edges
            batch_size = 1000
            for i in tqdm(range(0, len(edges), batch_size), desc="Importing edges"):
                batch = edges[i:i+batch_size]

                edge_data = []
                for source, target, attrs in batch:
                    edge_type = attrs.get('type', 'unknown')
                    edge_data.append({
                        'source': source,
                        'target': target,
                        'type': edge_type
                    })

                # Cypher query to create relationships
                query = """
                UNWIND $edges AS edge
                MATCH (source:Node {id: edge.source})
                MATCH (target:Node {id: edge.target})
                CREATE (source)-[r:DEPENDS_ON {type: edge.type}]->(target)
                """

                session.run(query, edges=edge_data)

            print(f"✅ Imported {len(edges)} edges")

    def create_specific_relationships(self, G):
        """Create specific relationship types based on edge types"""
        with self.driver.session() as session:
            # Get unique edge types
            edge_types = set()
            for _, _, data in G.edges(data=True):
                edge_types.add(data.get('type', 'unknown'))

            print(f"Found edge types: {edge_types}")

            # Create specific relationships for each type
            for edge_type in ['contains', 'imports', 'inherits', 'invokes']:
                if edge_type in edge_types:
                    query = f"""
                    MATCH (source:Node)-[r:DEPENDS_ON {{type: '{edge_type}'}}]->(target:Node)
                    CREATE (source)-[:{edge_type.upper()}]->(target)
                    """
                    result = session.run(query)
                    print(f"✅ Created {edge_type.upper()} relationships")

    def get_stats(self):
        """Get database statistics"""
        with self.driver.session() as session:
            node_count = session.run("MATCH (n) RETURN count(n) as count").single()['count']
            edge_count = session.run("MATCH ()-[r]->() RETURN count(r) as count").single()['count']

            # Get node type distribution
            type_dist = session.run("""
                MATCH (n:Node)
                RETURN n.type as type, count(n) as count
                ORDER BY count DESC
            """).data()

            # Get top nodes by degree
            top_nodes = session.run("""
                MATCH (n:Node)
                RETURN n.id as id, n.label as label, n.degree as degree
                ORDER BY n.degree DESC
                LIMIT 10
            """).data()

            return {
                'node_count': node_count,
                'edge_count': edge_count,
                'type_distribution': type_dist,
                'top_nodes': top_nodes
            }

def main():
    print("🚀 Starting Neo4j import process...")

    # Load the graph
    print("📊 Loading graph from pickle file...")
    with open("business-payout.pkl", "rb") as f:
        G = pickle.load(f)

    print(f"Graph loaded: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")

    # Connect to Neo4j
    importer = Neo4jImporter(URI, USERNAME, PASSWORD)

    try:
        # Clear existing data
        importer.clear_database()

        # Create indexes
        importer.create_indexes()

        # Import nodes
        importer.import_nodes(G)

        # Import edges
        importer.import_edges(G)

        # Create specific relationship types
        importer.create_specific_relationships(G)

        # Get and display stats
        stats = importer.get_stats()

        print("\n📈 Database Statistics:")
        print(f"Total nodes: {stats['node_count']}")
        print(f"Total edges: {stats['edge_count']}")

        print("\n📊 Node Type Distribution:")
        for item in stats['type_distribution'][:5]:
            print(f"  {item['type']}: {item['count']} nodes")

        print("\n🏆 Top Nodes by Degree:")
        for node in stats['top_nodes']:
            print(f"  {node['label']}: degree {node['degree']}")

        print("\n✅ Import completed successfully!")
        print("\n🌐 Neo4j Browser will be available at: http://localhost:7474")
        print("   Username: neo4j")
        print("   Password: neo4j123")

    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        importer.close()

if __name__ == "__main__":
    main()